﻿#pragma once

#include <optional>

#include <xy/math/math.h>
#include <xy/math/ndarray.h>

namespace xy
{

namespace math
{

template <std::size_t Dim, typename Real> class tensor : public ndarray<Dim, Real>
{
  protected:
    static_assert(Dim > 0, "Dim cannot be 0");
    using ndarray<Dim, Real>::m_shape;

    // 指标数组
    std::array<std::intptr_t, Dim> m_index;

  public:
    using ndarray<Dim, Real>::data;
    using ndarray<Dim, Real>::size;
    using ndarray<Dim, Real>::index;

    tensor(const tensor &) = default;
    tensor(tensor &&) = default;
    tensor &operator=(const tensor &) = default;
    tensor &operator=(tensor &&) = default;
    ~tensor() = default;

    /**
     * @brief Construct a new tensor object
     *
     * @param[in] base
     * @param[in] value
     */
    explicit tensor(std::size_t base, const Real &value)
        : ndarray<Dim, Real>(std::vector<Real>(power(base, Dim), value))
    {
        // 每个维度上的基数都相同
        m_shape.fill(base);
        m_index.fill(0);
    }

    /**
     * @brief Construct a new tensor object
     *
     * @param[in] base
     * @param[in] l
     */
    explicit tensor(std::size_t base, std::initializer_list<Real> l) : ndarray<Dim, Real>(l)
    {
        // 每个维度上的基数都相同
        m_shape.fill(base);
        m_index.fill(0);
    }

    /**
     * @brief Construct a new tensor object
     *
     * @param[in] base
     * @param[in] l
     */
    explicit tensor(std::size_t base, std::vector<Real> &&l) : ndarray<Dim, Real>(std::move(l))
    {
        // 每个维度上的基数都相同
        m_shape.fill(base);
        m_index.fill(0);
    }

    /**
     * @brief 得到基数
     *
     * @return std::size_t
     */
    std::size_t base() const noexcept
    {
        return m_shape[0];
    }

    /**
     * @brief 标准移动指标（从低位开始移动，包含进位操作），返回 false 说明移动后到达最后一个元素位置
     *
     * @return true
     * @return false
     */
    bool move_index()
    {
        std::size_t ind = Dim - 1;
        while (m_index[ind] == base() - 1 && ind > 0)
        {
            m_index[ind] = 0;
            ind--;
        }
        m_index[ind] = (m_index[ind] + 1) % base();
        return (ind != 0 || m_index[ind] != 0);
    }

    /**
     * @brief 移动指定维的指标，步长为 step
     *
     * @param[in] dim
     * @param[in] step
     */
    void move_index(std::size_t dim, int step = 1)
    {
        m_index[dim] = (m_index[dim] + step + base()) % base();
    }

    /**
     * @brief 得到当前指标
     *
     * @return const auto&
     */
    const auto &index() const noexcept
    {
        return m_index;
    }

    /**
     * @brief 得到当前指标
     *
     * @return auto&
     */
    auto &index() noexcept
    {
        return m_index;
    }

    /**
     * @brief 得到当前指标对应的值
     *
     * @param[in] step
     * @return Real
     */
    Real at_index(int step = 0) const
    {
        return data()[index(m_index) + step];
    }

    /**
     * @brief 得到当前指标对应的值
     *
     * @param[in] step
     * @return Real
     */
    Real &at_index(int step = 0)
    {
        return data()[index(m_index) + step];
    }

    /**
     * @brief 判断此维度上是上边界 1 还是下边界 -1 (均为 true)，或者不是边界 0 (false)
     *
     * @param[in] dim
     * @return int
     */
    int check_on_bound(std::size_t dim) const
    {
        if (m_index[dim] == 0)
            return -1;
        if (m_index[dim] == base() - 1)
            return 1;
        return 0;
    }

    /**
     * @brief 判断此维度上是靠近上边界 2,1 还是靠近下边界 -2,-1 (均为 true)，或者不靠近边界 0 (false)
     *
     * @param[in] dim
     * @return int
     */
    int check_near_bound(std::size_t dim) const
    {
        if (m_index[dim] == 0)
            return -2;
        if (m_index[dim] == 1)
            return -1;

        if (m_index[dim] == base() - 2)
            return 1;
        if (m_index[dim] == base() - 1)
            return 2;
        return 0;
    }

    /**
     * @brief 找到位于边界的维度
     *
     * @return int
     */
    std::optional<int> check_bounds() const
    {
        for (std::size_t d = 0; d < Dim; d++)
            if (check_on_bound(d))
                return d;
        return std::nullopt;
    }

// 关于数的四则运算 + - * /
#define TENSOR_BINARY_OP(OpName, Op)                                                                                   \
    tensor OpName(Real r) const                                                                                        \
    {                                                                                                                  \
        tensor t = *this;                                                                                              \
        std::size_t n = size();                                                                                        \
        for (std::size_t i = 0; i < n; i++)                                                                            \
            t.data()[i] = t.data()[i] Op r;                                                                            \
        return t;                                                                                                      \
    }

    TENSOR_BINARY_OP(operator+, +)
    TENSOR_BINARY_OP(operator-, -)
    TENSOR_BINARY_OP(operator*, *)
    TENSOR_BINARY_OP(operator/, /)
#undef TENSOR_BINARY_OP

// 关于数的四则运算 += -= *= /=
#define TENSOR_BINARY_OP(OpName, Op)                                                                                   \
    tensor &OpName(Real r)                                                                                             \
    {                                                                                                                  \
        std::size_t n = size();                                                                                        \
        for (std::size_t i = 0; i < n; i++)                                                                            \
            data()[i] Op r;                                                                                            \
        return *this;                                                                                                  \
    }

    TENSOR_BINARY_OP(operator+=, +=)
    TENSOR_BINARY_OP(operator-=, -=)
    TENSOR_BINARY_OP(operator*=, *=)
    TENSOR_BINARY_OP(operator/=, /=)
#undef TENSOR_BINARY_OP

// 关于张量的四则运算 + - * /
#define TENSOR_BINARY_OP(OpName, Op)                                                                                   \
    tensor OpName(const tensor &rhs) const                                                                             \
    {                                                                                                                  \
        tensor t = *this;                                                                                              \
        std::size_t n = size();                                                                                        \
        for (std::size_t i = 0; i < n; i++)                                                                            \
            t.data()[i] = t.data()[i] Op rhs.data()[i];                                                                \
        return t;                                                                                                      \
    }

    TENSOR_BINARY_OP(operator+, +)
    TENSOR_BINARY_OP(operator-, -)
#undef TENSOR_BINARY_OP

// 关于张量的四则运算 += -= *= /=
#define TENSOR_BINARY_OP(OpName, Op)                                                                                   \
    tensor &OpName(const tensor &rhs)                                                                                  \
    {                                                                                                                  \
        std::size_t n = size();                                                                                        \
        for (std::size_t i = 0; i < n; i++)                                                                            \
            data()[i] Op rhs.data()[i];                                                                                \
        return *this;                                                                                                  \
    }

    TENSOR_BINARY_OP(operator+=, +=)
    TENSOR_BINARY_OP(operator-=, -=)
#undef TENSOR_BINARY_OP

    /**
     * @brief 取反
     *
     * @return tensor
     */
    tensor operator-() const
    {
        return *this * -1;
    }

    /**
     * @brief 张量点积
     *
     * @param[in] rhs
     * @return Real
     */
    Real operator*(const tensor &rhs)
    {
        Real ret = 0;
        std::size_t n = size();
        for (std::size_t i = 0; i < n; i++)
            ret += data()[i] * rhs.data()[i];
        return ret;
    }

    /**
     * @brief 以 i 为起点，返回 D 阶 Tensor
     *
     * @tparam D
     * @param[in] i
     * @return tensor<D, Real>
     */
    template <std::size_t D> tensor<D, Real> component(std::size_t i) const
    {
        // 步长
        auto start = power<int>(base(), D) * i;
        tensor<D, Real> t(base(), 0);
        std::size_t n = t.size();
        for (std::size_t j = 0; j < n; j++)
            t.data()[j] = data()[start + j];
        return t;
    }

    /**
     * @brief 张量乘法：Dim - D 阶张量 = Dim 阶张量 * D 阶张量
     *
     * @tparam D
     * @param[in] rhs
     * @return tensor<Dim - D, Real>
     */
    template <std::size_t D> tensor<Dim - D, Real> operator*(const tensor<D, Real> &rhs) const
    {
        tensor<Dim - D, Real> t(base(), 0);
        std::size_t n = t.size();
        for (std::size_t i = 0; i < n; ++i)
            t.data()[i] = component<D>(i) * rhs;
        return t;
    }

    /**
     * @brief 获得包围当前值的所有值
     *
     * @return tensor<Dim, Real>
     */
    tensor<Dim, Real> cell() const
    {
        // 检测对应点是否在边界上，只要在边界上就报错
        if (check_bounds())
            throw std::out_of_range("cell out of range");

        // 模拟指标，每个分量 -1，进而移动到左下角
        std::array<std::intptr_t, Dim> corner = m_index;
        for (std::size_t i = 0; i < Dim; i++)
            corner[i]--;

        std::array<std::intptr_t, Dim> added;
        tensor<Dim, Real> val(3, 0);
        do
        {
            for (std::size_t d = 0; d < Dim; d++)
                added[d] = corner[d] + val.index()[d];
            val.at_index() = data()[index(added)];
        } while (val.move_index());
        return val;
    }
};

/**
 * @brief 一维张量构建多维张量
 *
 * @tparam N
 * @tparam T
 * @param[in] t
 * @return tensor<Dim, T>
 */
template <std::size_t Dim, class T> inline tensor<Dim, T> power(const tensor<1, T> &t)
{
    std::size_t n = t.size();
    tensor<Dim, T> ret(n, 0);
    for (std::size_t i = 0; i < n; i++)
        ret.data()[i] = t.data()[i];

    tensor<Dim, T> record(n, 0);
    for (std::size_t d = 1; d < Dim; d++)
    {
        auto step = std::pow(n, d);
        for (std::size_t i = 0; i < step; i++)
            for (std::size_t j = 0; j < n; j++)
                record.data()[i * n + j] = ret.data()[i] * t.data()[j];

        for (std::size_t i = 0; i < step * n; i++)
            ret.data()[i] = record.data()[i];
    }
    return ret;
}

} // namespace math

} // namespace xy
